import os
import cv2

class loader:
    def __init__(self, class_name):
        self.class_name = class_name
        self.class_dir = 'E:/code/vot2014/' + self.class_name + '/'

    def frame_generator(self, v_dir=None):
        frame_seq = []
        gt_seq = []
        if v_dir != None and self.class_name == None:
            cap = cv2.VideoCapture(v_dir)
            while cap.isOpened():
                ret, frame = cap.read()
                if ret == True:
                    frame_seq.append(frame)
            cap.release()
        elif v_dir == None and self.class_name != None:
            i = 1
            file_name = ''
            while True:
                file_index = str(i)
                file_name = ''
                for j in range(0, 8 - len(file_index)):
                    file_name = file_name + '0'
                file_name = file_name + file_index + '.jpg'
                if os.path.exists(self.class_dir + file_name):
                    img = cv2.imread(self.class_dir + file_name)
                    frame_seq.append(img)
                    i += 1
                else:
                    break
        
        return frame_seq

    def gt_generator(self):
        gt_seq = []
        with open(self.class_dir + 'groundtruth.txt') as fp:
            while True:
                strgt = fp.readline()
                if len(strgt) == 0:
                    break
                data = strgt.split(',')
                for i in range(0, len(data)):
                    data[i] = int(float(data[i]))
                x, y = min(data[0], data[4]), min(data[1], data[5])
                w, h = abs(data[0] - data[4]), abs(data[1] - data[5])
                gt_seq.append([x, y, w, h])
        return gt_seq